import os.path
import cPickle

from mercurial.commands import commit
from mercurial.commands import import_
from mercurial.commands import merge
from mercurial.commands import resolve
from mercurial.commands import revert
from mercurial.commands import update
from mercurial.patch import extract


BASE_DIR = 'ppatches'
PATCH_DIR = 'patches'
REVMAP_FILE = 'revmap.pkl'


class PatchData(object):
    
    def __init__(self, ui, filename):
        with open(filename) as fp:
            data = extract(ui, fp)
        os.unlink(data[0]) # tmpname

        self.filename = filename
        self.message = data[1]
        self.user = data[2]
        self.date = data[3]
        self.branch = data[4]
        self.node = data[5]
        self.p1 = data[6]
        self.p2 = data[7]


class RevMap(dict):

    def __init__(self, repo):
        self.repo = repo

        path = os.path.join(repo.path, BASE_DIR, REVMAP_FILE)
        try:
            with open(path, 'r') as fp:
                d = cPickle.load(fp)
            for key, value in d.items():
                self[key] = value
        except IOError:
            pass

    def __getitem__(self, key):
        if self.has_key(key):
            return super(RevMap, self).get(key)
        return key

    def save(self):
        path = os.path.join(self.repo.path, BASE_DIR, REVMAP_FILE)
        with open(path, 'w') as fp:
            cPickle.dump(dict(self), fp)


def revrange(rev):
    revs = rev.split(':', 1)
    revs = map(int, revs)
    if len(revs) > 1:
        revs = range(revs[0], revs[1] + 1)
    else:
        revs = [revs[0]]
    return revs


def pimport_mainloop(ui, repo, patches, revmap):
    for patch in patches:
        pdata = PatchData(ui, patch)
        
        # always start from the correct parent
        update(ui, repo, revmap[pdata.p1])

        if pdata.p2 is None:
            # branch when needed (no support yet for named branches)
            import_(ui, repo, patch, base='', strip=1)
        else:
            # merge using the patch to resolve conflicts
            merge(ui, repo, rev=revmap[pdata.p2])
            resolve(ui, repo, all=True)
            revert(ui, repo, date='', all=True, rev=revmap[pdata.p1], no_backup=True)
            import_(ui, repo, pdata.filename, base='', strip=1, force=True, no_commit=True)
            repo.commit(pdata.message, pdata.user, pdata.date)

        # update the revision map
        if repo['tip'].hex() != pdata.node:
            revmap[pdata.node] = repo['tip'].hex()


def pimport(ui, repo, rev):
    revs = revrange(rev)
    patch_filepath = os.path.join(repo.path, BASE_DIR, PATCH_DIR, '%i.diff')
    patches = map(lambda x: patch_filepath % x, revs)

    revmap = RevMap(repo)
    try:
        pimport_mainloop(ui, repo, patches, revmap)
    finally:
        if revmap:
            revmap.save()


cmdtable = {
    "pimport": (pimport, [], "REV")
}
